
/*
 * mbuf.c
 *		内存缓冲区操作。
 *
 * Copyright (c) 2005 Marko Kreen
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *	  notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *	  notice, this list of conditions and the following disclaimer in the
 *	  documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 *
 * contrib/pgcrypto/mbuf.c
 */

#include "postgres.h"

#include "mbuf.h"
#include "px.h"

#define STEP  (16*1024)

struct MBuf
{
	uint8	   *data;
	uint8	   *data_end;
	uint8	   *read_pos;
	uint8	   *buf_end;
	bool		no_write;
	bool		own_data;
};

int mbuf_avail(MBuf *fc_mbuf)
{
	return fc_mbuf->data_end - fc_mbuf->read_pos;
}

int mbuf_size(MBuf *fc_mbuf)
{
	return fc_mbuf->data_end - fc_mbuf->data;
}

int mbuf_tell(MBuf *fc_mbuf)
{
	return fc_mbuf->read_pos - fc_mbuf->data;
}

int mbuf_free(MBuf *fc_mbuf)
{
	if (fc_mbuf->own_data)
	{
		px_memset(fc_mbuf->data, 0, fc_mbuf->buf_end - fc_mbuf->data);
		pfree(fc_mbuf->data);
	}
	pfree(fc_mbuf);
	return 0;
}

static void fc_prepare_room(MBuf *fc_mbuf, int fc_block_len)
{
	uint8	   *fc_newbuf;
	unsigned	fc_newlen;

	if (fc_mbuf->data_end + fc_block_len <= fc_mbuf->buf_end)
		return;

	fc_newlen = (fc_mbuf->buf_end - fc_mbuf->data)
		+ ((fc_block_len + STEP + STEP - 1) & -STEP);

	fc_newbuf = repalloc(fc_mbuf->data, fc_newlen);

	fc_mbuf->buf_end = fc_newbuf + fc_newlen;
	fc_mbuf->data_end = fc_newbuf + (fc_mbuf->data_end - fc_mbuf->data);
	fc_mbuf->read_pos = fc_newbuf + (fc_mbuf->read_pos - fc_mbuf->data);
	fc_mbuf->data = fc_newbuf;
}

int mbuf_append(MBuf *fc_dst, const uint8 *fc_buf, int fc_len)
{
	if (fc_dst->no_write)
	{
		px_debug("mbuf_append: no_write");
		return PXE_BUG;
	}

	fc_prepare_room(fc_dst, fc_len);

	memcpy(fc_dst->data_end, fc_buf, fc_len);
	fc_dst->data_end += fc_len;

	return 0;
}

MBuf *
mbuf_create(int fc_len)
{
	MBuf	   *fc_mbuf;

	if (!fc_len)
		fc_len = 8192;

	fc_mbuf = palloc(sizeof *fc_mbuf);
	fc_mbuf->data = palloc(fc_len);
	fc_mbuf->buf_end = fc_mbuf->data + fc_len;
	fc_mbuf->data_end = fc_mbuf->data;
	fc_mbuf->read_pos = fc_mbuf->data;

	fc_mbuf->no_write = false;
	fc_mbuf->own_data = true;

	return fc_mbuf;
}

MBuf *
mbuf_create_from_data(uint8 *fc_data, int fc_len)
{
	MBuf	   *fc_mbuf;

	fc_mbuf = palloc(sizeof *fc_mbuf);
	fc_mbuf->data = (uint8 *) fc_data;
	fc_mbuf->buf_end = fc_mbuf->data + fc_len;
	fc_mbuf->data_end = fc_mbuf->data + fc_len;
	fc_mbuf->read_pos = fc_mbuf->data;

	fc_mbuf->no_write = true;
	fc_mbuf->own_data = false;

	return fc_mbuf;
}


int mbuf_grab(MBuf *fc_mbuf, int fc_len, uint8 **fc_data_p)
{
	if (fc_len > mbuf_avail(fc_mbuf))
		fc_len = mbuf_avail(fc_mbuf);

	fc_mbuf->no_write = true;

	*fc_data_p = fc_mbuf->read_pos;
	fc_mbuf->read_pos += fc_len;
	return fc_len;
}

int mbuf_rewind(MBuf *fc_mbuf)
{
	fc_mbuf->read_pos = fc_mbuf->data;
	return 0;
}

int mbuf_steal_data(MBuf *fc_mbuf, uint8 **fc_data_p)
{
	int			fc_len = mbuf_size(fc_mbuf);

	fc_mbuf->no_write = true;
	fc_mbuf->own_data = false;

	*fc_data_p = fc_mbuf->data;

	fc_mbuf->data = fc_mbuf->data_end = fc_mbuf->read_pos = fc_mbuf->buf_end = NULL;

	return fc_len;
}

/*
 * PullFilter
 */

struct PullFilter
{
	PullFilter *src;
	const PullFilterOps *op;
	int			buflen;
	uint8	   *buf;
	int			pos;
	void	   *priv;
};

int pullf_create(PullFilter **fc_pf_p, const PullFilterOps *fc_op, void *fc_init_arg, PullFilter *fc_src)
{
	PullFilter *fc_pf;
	void	   *fc_priv;
	int			fc_res;

	if (fc_op->init != NULL)
	{
		fc_res = fc_op->init(&fc_priv, fc_init_arg, fc_src);
		if (fc_res < 0)
			return fc_res;
	}
	else
	{
		fc_priv = fc_init_arg;
		fc_res = 0;
	}

	fc_pf = palloc0(sizeof(*fc_pf));
	fc_pf->buflen = fc_res;
	fc_pf->op = fc_op;
	fc_pf->priv = fc_priv;
	fc_pf->src = fc_src;
	if (fc_pf->buflen > 0)
	{
		fc_pf->buf = palloc(fc_pf->buflen);
		fc_pf->pos = 0;
	}
	else
	{
		fc_pf->buf = NULL;
		fc_pf->pos = 0;
	}
	*fc_pf_p = fc_pf;
	return 0;
}

void pullf_free(PullFilter *fc_pf)
{
	if (fc_pf->op->free)
		fc_pf->op->free(fc_pf->priv);

	if (fc_pf->buf)
	{
		px_memset(fc_pf->buf, 0, fc_pf->buflen);
		pfree(fc_pf->buf);
	}

	px_memset(fc_pf, 0, sizeof(*fc_pf));
	pfree(fc_pf);
}

/* 可能返回的数据显示少于请求的数据，0表示文件结束 */
int pullf_read(PullFilter *fc_pf, int fc_len, uint8 **fc_data_p)
{
	int			fc_res;

	if (fc_pf->op->pull)
	{
		if (fc_pf->buflen && fc_len > fc_pf->buflen)
			fc_len = fc_pf->buflen;
		fc_res = fc_pf->op->pull(fc_pf->priv, fc_pf->src, fc_len, fc_data_p,
						   fc_pf->buf, fc_pf->buflen);
	}
	else
		fc_res = pullf_read(fc_pf->src, fc_len, fc_data_p);
	return fc_res;
}

int pullf_read_max(PullFilter *fc_pf, int fc_len, uint8 **fc_data_p, uint8 *fc_tmpbuf)
{
	int			fc_res,
				fc_total;
	uint8	   *fc_tmp;

	fc_res = pullf_read(fc_pf, fc_len, fc_data_p);
	if (fc_res <= 0 || fc_res == fc_len)
		return fc_res;

	/* 读取的长度较短，使用tmpbuf */
	memcpy(fc_tmpbuf, *fc_data_p, fc_res);
	*fc_data_p = fc_tmpbuf;
	fc_len -= fc_res;
	fc_total = fc_res;

	while (fc_len > 0)
	{
		fc_res = pullf_read(fc_pf, fc_len, &fc_tmp);
		if (fc_res < 0)
		{
			/* 所以调用者必须仅在成功时清除 */
			px_memset(fc_tmpbuf, 0, fc_total);
			return fc_res;
		}
		if (fc_res == 0)
			break;
		memcpy(fc_tmpbuf + fc_total, fc_tmp, fc_res);
		fc_total += fc_res;
		fc_len -= fc_res;
	}
	return fc_total;
}

/*
 * 调用者想要确切的len字节，不用担心引用
 */
int pullf_read_fixed(PullFilter *fc_src, int fc_len, uint8 *fc_dst)
{
	int			fc_res;
	uint8	   *fc_p;

	fc_res = pullf_read_max(fc_src, fc_len, &fc_p, fc_dst);
	if (fc_res < 0)
		return fc_res;
	if (fc_res != fc_len)
	{
		px_debug("pullf_read_fixed: need=%d got=%d", fc_len, fc_res);
		return PXE_PGP_CORRUPT_DATA;
	}
	if (fc_p != fc_dst)
		memcpy(fc_dst, fc_p, fc_len);
	return 0;
}

/*
 * 从MBuf读取
 */
static int fc_pull_from_mbuf(void *fc_arg, PullFilter *fc_src, int fc_len,
			   uint8 **fc_data_p, uint8 *fc_buf, int fc_buflen)
{
	MBuf	   *fc_mbuf = fc_arg;

	return mbuf_grab(fc_mbuf, fc_len, fc_data_p);
}

static const struct PullFilterOps mbuf_reader = {
	NULL, fc_pull_from_mbuf, NULL
};

int pullf_create_mbuf_reader(PullFilter **fc_mp_p, MBuf *fc_src)
{
	return pullf_create(fc_mp_p, &mbuf_reader, fc_src, NULL);
}


/*
 * PushFilter
 */

struct PushFilter
{
	PushFilter *next;
	const PushFilterOps *op;
	int			block_size;
	uint8	   *buf;
	int			pos;
	void	   *priv;
};

int pushf_create(PushFilter **fc_mp_p, const PushFilterOps *fc_op, void *fc_init_arg, PushFilter *fc_next)
{
	PushFilter *fc_mp;
	void	   *fc_priv;
	int			fc_res;

	if (fc_op->init != NULL)
	{
		fc_res = fc_op->init(fc_next, fc_init_arg, &fc_priv);
		if (fc_res < 0)
			return fc_res;
	}
	else
	{
		fc_priv = fc_init_arg;
		fc_res = 0;
	}

	fc_mp = palloc0(sizeof(*fc_mp));
	fc_mp->block_size = fc_res;
	fc_mp->op = fc_op;
	fc_mp->priv = fc_priv;
	fc_mp->next = fc_next;
	if (fc_mp->block_size > 0)
	{
		fc_mp->buf = palloc(fc_mp->block_size);
		fc_mp->pos = 0;
	}
	else
	{
		fc_mp->buf = NULL;
		fc_mp->pos = 0;
	}
	*fc_mp_p = fc_mp;
	return 0;
}

void pushf_free(PushFilter *fc_mp)
{
	if (fc_mp->op->free)
		fc_mp->op->free(fc_mp->priv);

	if (fc_mp->buf)
	{
		px_memset(fc_mp->buf, 0, fc_mp->block_size);
		pfree(fc_mp->buf);
	}

	px_memset(fc_mp, 0, sizeof(*fc_mp));
	pfree(fc_mp);
}

void pushf_free_all(PushFilter *fc_mp)
{
	PushFilter *fc_tmp;

	while (fc_mp)
	{
		fc_tmp = fc_mp->next;
		pushf_free(fc_mp);
		fc_mp = fc_tmp;
	}
}

static int fc_wrap_process(PushFilter *fc_mp, const uint8 *fc_data, int fc_len)
{
	int			fc_res;

	if (fc_mp->op->push != NULL)
		fc_res = fc_mp->op->push(fc_mp->next, fc_mp->priv, fc_data, fc_len);
	else
		fc_res = pushf_write(fc_mp->next, fc_data, fc_len);
	if (fc_res > 0)
		return PXE_BUG;
	return fc_res;
}

/* 消耗所有数据，成功时返回len */
int pushf_write(PushFilter *fc_mp, const uint8 *fc_data, int fc_len)
{
	int			fc_need,
				fc_res;

	/*
	 * 不进行缓冲
	 */
	if (fc_mp->block_size <= 0)
		return fc_wrap_process(fc_mp, fc_data, fc_len);

	/*
	 * 尝试清空缓冲区
	 */
	fc_need = fc_mp->block_size - fc_mp->pos;
	if (fc_need > 0)
	{
		if (fc_len < fc_need)
		{
			memcpy(fc_mp->buf + fc_mp->pos, fc_data, fc_len);
			fc_mp->pos += fc_len;
			return 0;
		}
		memcpy(fc_mp->buf + fc_mp->pos, fc_data, fc_need);
		fc_len -= fc_need;
		fc_data += fc_need;
	}

	/*
	 * 缓冲区满，处理
	 */
	fc_res = fc_wrap_process(fc_mp, fc_mp->buf, fc_mp->block_size);
	if (fc_res < 0)
		return fc_res;
	fc_mp->pos = 0;

	/*
	 * 现在直接从数据中处理
	 */
	while (fc_len > 0)
	{
		if (fc_len > fc_mp->block_size)
		{
			fc_res = fc_wrap_process(fc_mp, fc_data, fc_mp->block_size);
			if (fc_res < 0)
				return fc_res;
			fc_data += fc_mp->block_size;
			fc_len -= fc_mp->block_size;
		}
		else
		{
			memcpy(fc_mp->buf, fc_data, fc_len);
			fc_mp->pos += fc_len;
			break;
		}
	}
	return 0;
}

int pushf_flush(PushFilter *fc_mp)
{
	int			fc_res;

	while (fc_mp)
	{
		if (fc_mp->block_size > 0)
		{
			fc_res = fc_wrap_process(fc_mp, fc_mp->buf, fc_mp->pos);
			if (fc_res < 0)
				return fc_res;
		}

		if (fc_mp->op->flush)
		{
			fc_res = fc_mp->op->flush(fc_mp->next, fc_mp->priv);
			if (fc_res < 0)
				return fc_res;
		}

		fc_mp = fc_mp->next;
	}
	return 0;
}


/*
 * 写入MBuf
 */
static int fc_push_into_mbuf(PushFilter *fc_next, void *fc_arg, const uint8 *fc_data, int fc_len)
{
	int			fc_res = 0;
	MBuf	   *fc_mbuf = fc_arg;

	if (fc_len > 0)
		fc_res = mbuf_append(fc_mbuf, fc_data, fc_len);
	return fc_res < 0 ? fc_res : 0;
}

static const struct PushFilterOps mbuf_filter = {
	NULL, fc_push_into_mbuf, NULL, NULL
};

int pushf_create_mbuf_writer(PushFilter **fc_res, MBuf *fc_dst)
{
	return pushf_create(fc_res, &mbuf_filter, fc_dst, NULL);
}
